import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np
import os
from sklearn.metrics import confusion_matrix, roc_curve, auc
from wordcloud import WordCloud

class ClassificationVisualizer:
    def __init__(self):
        plt.style.use('seaborn-v0_8-darkgrid')
        self.colors = plt.cm.Set3(np.linspace(0, 1, 10))
        plt.rcParams['font.sans-serif'] = ['PingFang HK', 'Arial']
        plt.rcParams['axes.unicode_minus'] = False
    
    def plot_confusion_matrix(self, y_true, y_pred, labels=None):
        """绘制混淆矩阵热力图"""
        plt.figure(figsize=(10, 8))
        cm = confusion_matrix(y_true, y_pred)
        sns.heatmap(cm, annot=True, fmt='d', cmap='Blues',
                    xticklabels=labels, yticklabels=labels, square=True)
        plt.title('分类混淆矩阵')
        plt.xlabel('预测类别')
        plt.ylabel('真实类别')
        plt.tight_layout()
        return plt.gcf()
    
    def plot_accuracy_bar(self, class_accuracies, labels=None):
        """绘制分类准确率柱状图"""
        plt.figure(figsize=(10, 6))
        x = range(len(class_accuracies))
        plt.bar(x, class_accuracies, color=self.colors)
        if labels:
            plt.xticks(x, labels, rotation=45)
        plt.title('各类别分类准确率')
        plt.xlabel('类别')
        plt.ylabel('准确率')
        plt.ylim(0, 1.0)
        plt.tight_layout()
        return plt.gcf()
    
    def plot_roc_curve(self, y_true, y_prob):
        """绘制ROC曲线"""
        plt.figure(figsize=(8, 8))
        fpr, tpr, _ = roc_curve(y_true, y_prob)
        roc_auc = auc(fpr, tpr)
        
        plt.plot(fpr, tpr, color='#2ecc71',
                 lw=2, label=f'ROC曲线 (AUC = {roc_auc:.2f})')
        plt.plot([0, 1], [0, 1], color='#e74c3c', lw=2, linestyle='--')
        plt.xlim([0.0, 1.0])
        plt.ylim([0.0, 1.05])
        plt.xlabel('假阳性率')
        plt.ylabel('真阳性率')
        plt.title('ROC曲线')
        plt.legend(loc='lower right')
        plt.grid(True)
        plt.tight_layout()
        return plt.gcf()
    
    def plot_wordcloud(self, text_data):
        """生成词云图"""
        plt.figure(figsize=(12, 8))
        try:
            # 尝试使用系统字体
            font_path = '/System/Library/Fonts/PingFang.ttc'
            if not os.path.exists(font_path):
                # 如果PingFang字体不存在，尝试使用Arial字体
                font_path = '/Library/Fonts/Arial Unicode.ttf'
                if not os.path.exists(font_path):
                    # 如果Arial字体也不存在，使用系统默认字体
                    font_path = None
            
            wordcloud = WordCloud(width=1200, height=800,
                                background_color='white',
                                font_path=font_path,
                                max_words=100,
                                colormap='Set3',
                                prefer_horizontal=0.7)
            wordcloud.generate(text_data)
        except Exception as e:
            print(f"生成词云图时出错：{e}")
            return None
        plt.imshow(wordcloud, interpolation='bilinear')
        plt.axis('off')
        plt.title('文档关键词云图')
        plt.tight_layout()
        return plt.gcf()
    
    def save_plots(self, plots, filenames):
        """保存所有可视化图表"""
        for plot, filename in zip(plots, filenames):
            plot.savefig(filename, dpi=300, bbox_inches='tight')
            plt.close(plot)